Speed up clearing downstream tasks on large Dags - #71203

from __future__ import annotations

import time

import pendulum

from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.sdk import DAG
from airflow.serialization.serialized_objects import DagSerialization


def _serialize(dag):
    return DagSerialization.from_dict(DagSerialization.to_dict(dag))


def build_chain(n: int):
    """t0 -> t1 -> ... -> tn. Every task reaches every later one."""
    with DAG("bench", start_date=pendulum.datetime(2024, 1, 1), schedule=None) as dag:
        previous = None
        for i in range(n):
            task = EmptyOperator(task_id=f"t{i:04d}")
            if previous is not None:
                previous >> task
            previous = task
    return _serialize(dag)


def build_parallel_chains(n: int, streams: int = 10):
    """`streams` independent chains side by side."""
    with DAG("bench", start_date=pendulum.datetime(2024, 1, 1), schedule=None) as dag:
        tails: dict[int, object] = {}
        for i in range(n):
            task = EmptyOperator(task_id=f"t{i:04d}")
            stream = i % streams
            if (previous := tails.get(stream)) is not None:
                previous >> task
            tails[stream] = task
    return _serialize(dag)


def build_layered(n: int, width: int = 20):
    """Layers of `width` tasks, each fully connected to the next layer."""
    with DAG("bench", start_date=pendulum.datetime(2024, 1, 1), schedule=None) as dag:
        tasks = [EmptyOperator(task_id=f"t{i:04d}") for i in range(n)]
        for start in range(0, n - width, width):
            for upstream in tasks[start : start + width]:
                for downstream in tasks[start + width : start + 2 * width]:
                    upstream >> downstream
    return _serialize(dag)


def build_fanout(n: int):
    """One root feeding every other task -- the shallow extreme."""
    with DAG("bench", start_date=pendulum.datetime(2024, 1, 1), schedule=None) as dag:
        root = EmptyOperator(task_id="root")
        for i in range(n - 1):
            root >> EmptyOperator(task_id=f"t{i:04d}")
    return _serialize(dag)


SHAPES = [
    ("chain", build_chain),
    ("parallel chains x10", build_parallel_chains),
    ("layered (w=20)", build_layered),
    ("fan-out", build_fanout),
]


def timed(dag, task_ids) -> float:
    started = time.monotonic()
    dag.partial_subset(task_ids=task_ids, include_downstream=True, include_upstream=False)
    return time.monotonic() - started


print(f"{'shape':<16} {'tasks':>6} {'relatives':>11} {'partial_subset':>16}")
for label, build in SHAPES:
    for n in (100, 200, 400):
        dag = build(n)
        task_ids = {t.task_id for t in dag.tasks}
        relatives = sum(len(t.get_flat_relatives(upstream=False, depth=None)) for t in dag.tasks)
        best = min(timed(dag, task_ids) for _ in range(3))
        print(f"{label:<16} {n:>6} {relatives:>11,} {best:>15.3f}s")
    print()